#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
DICOM to NIfTI Converter GUI
基于PyQt5的DICOM文件到NIfTI文件转换工具
"""

import os
import sys
import logging
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                            QPushButton, QLabel, QLineEdit, QFileDialog, QMessageBox,
                            QGroupBox, QRadioButton, QButtonGroup, QTextEdit, QScrollArea)
from PyQt5.QtCore import QThread, pyqtSignal, QTimer

# 导入dicom2nifti库
try:
    import dicom2nifti
    DICOM2NIFTI_AVAILABLE = True
except ImportError:
    DICOM2NIFTI_AVAILABLE = False
    logging.error("未找到dicom2nifti库，请使用pip install dicom2nifti进行安装")


class QTextEditLogger(logging.Handler):
    """
    自定义日志处理器，用于将日志信息输出到QTextEdit控件
    """
    def __init__(self, text_edit):
        super().__init__()
        self.text_edit = text_edit
        self.text_edit.setReadOnly(True)
    
    def emit(self, record):
        """
        重写emit方法，将日志记录发送到QTextEdit控件
        
        Args:
            record: 日志记录对象
        """
        msg = self.format(record)
        self.text_edit.append(msg)
        # 自动滚动到底部
        self.text_edit.verticalScrollBar().setValue(
            self.text_edit.verticalScrollBar().maximum()
        )


class ConversionWorker(QThread):
    """
    转换工作线程，用于在后台执行DICOM到NIfTI的转换操作
    """
    finished = pyqtSignal(bool, str)
    
    def __init__(self, conversion_type, dicom_path, output_path, reorient=True, compression=True, dicom_files=None):
        """
        初始化转换工作线程
        
        Args:
            conversion_type: 转换类型 ('directory', 'series' 或 'files')
            dicom_path: DICOM文件目录路径或DICOM系列路径
            output_path: 输出路径（目录或文件）
            reorient: 是否重定向NIfTI文件
            compression: 是否压缩NIfTI文件（仅对目录转换有效）
            dicom_files: 要转换的DICOM文件列表（当conversion_type为'files'时使用）
        """
        super().__init__()
        self.conversion_type = conversion_type
        self.dicom_path = dicom_path  # 重命名为更通用的名称
        self.output_path = output_path
        self.reorient = reorient
        self.compression = compression
        self.dicom_files = dicom_files or []  # 存储要转换的文件列表
    
    def run(self):
        """
        执行转换操作，直接进行DICOM到NIfTI的转换，不显示进度条
        """
        try:
            if self.conversion_type == 'directory':
                # 直接执行DICOM目录到NIfTI的转换
                dicom2nifti.convert_directory(
                    self.dicom_path, 
                    self.output_path, 
                    compression=self.compression, 
                    reorient=self.reorient
                )
                
            elif self.conversion_type == 'series':
                # 直接执行DICOM系列到NIfTI的转换
                dicom2nifti.dicom_series_to_nifti(
                    self.dicom_path, 
                    self.output_path, 
                    reorient_nifti=self.reorient
                )
                
            elif self.conversion_type == 'files':
                # 处理选中的DICOM文件
                import tempfile
                import shutil
                
                with tempfile.TemporaryDirectory() as temp_dir:
                    # 复制选中的文件到临时目录
                    for file_path in self.dicom_files:
                        if os.path.exists(file_path):
                            shutil.copy2(file_path, temp_dir)
                    
                    # 执行转换
                    dicom2nifti.dicom_series_to_nifti(
                        temp_dir, 
                        self.output_path, 
                        reorient_nifti=self.reorient
                    )
            
            self.finished.emit(True, "转换成功完成！")
            
        except Exception as e:
            logging.error(f"转换过程中出错: {str(e)}")
            self.finished.emit(False, f"转换失败: {str(e)}")


class Dicom2NiftiConverter(QMainWindow):
    """
    DICOM到NIfTI转换工具的主窗口类
    """
    def __init__(self):
        """
        初始化主窗口
        """
        super().__init__()
        self.init_ui()
        
        # 初始化日志相关变量，但不立即设置日志处理器
        self.last_log_size = 0
        self.log_timer = QTimer(self)
        self.log_timer.timeout.connect(self.refresh_log_file)
        self.log_timer.setSingleShot(True)  # 初始设置为单次触发模式
        
        # 初始化日志标志
        self.logging_initialized = False
    
    def setup_logging(self):
        """
        设置日志系统，包括文件日志和UI日志
        """
        # 清除现有的日志处理器
        root_logger = logging.getLogger()
        root_logger.handlers.clear()
        
        # 设置日志格式
        log_format = '%(asctime)s - %(levelname)s - %(message)s'
        formatter = logging.Formatter(log_format)
        
        # 添加文件日志处理器
        file_handler = logging.FileHandler("dicom2nifti_converter.log", mode='w')  # 使用'w'模式覆盖现有日志
        file_handler.setFormatter(formatter)
        root_logger.addHandler(file_handler)
        
        # 添加控制台日志处理器
        stream_handler = logging.StreamHandler()
        stream_handler.setFormatter(formatter)
        root_logger.addHandler(stream_handler)
        
        # 添加UI日志处理器
        ui_handler = QTextEditLogger(self.log_text_edit)
        ui_handler.setFormatter(formatter)
        root_logger.addHandler(ui_handler)
        
        # 设置日志级别
        root_logger.setLevel(logging.INFO)
        
        # 标记日志系统已初始化
        self.logging_initialized = True
    
    def refresh_log_file(self):
        """
        实时刷新日志文件内容到界面
        """
        log_file_path = "dicom2nifti_converter.log"
        try:
            if os.path.exists(log_file_path):
                # 获取文件当前大小
                current_size = os.path.getsize(log_file_path)
                
                # 如果文件大小有变化，读取新增内容
                if current_size > self.last_log_size:
                    with open(log_file_path, 'r', encoding='utf-8', errors='replace') as f:
                        f.seek(self.last_log_size)
                        new_content = f.read()
                        
                        if new_content:
                            # 按行分割并添加到日志显示
                            for line in new_content.strip().split('\n'):
                                if line.strip():
                                    # 避免重复添加
                                    if not self.log_text_edit.toPlainText().endswith(line):
                                        self.log_text_edit.append(line)
                                        # 自动滚动到底部
                                        self.log_text_edit.verticalScrollBar().setValue(
                                            self.log_text_edit.verticalScrollBar().maximum()
                                        )
                    
                    # 更新最后读取的文件大小
                    self.last_log_size = current_size
        except Exception as e:
            # 避免刷新日志时的错误影响主程序运行
            pass
        

    
    def center_on_screen(self):
        """
        将窗口居中显示在屏幕上
        """
        # 获取屏幕的几何信息
        screen_geometry = QApplication.desktop().screenGeometry()
        # 获取窗口的几何信息
        window_geometry = self.frameGeometry()
        # 计算窗口的中心位置（基于屏幕中心）
        center_point = screen_geometry.center()
        # 将窗口的中心与计算出的中心位置对齐
        window_geometry.moveCenter(center_point)
        # 将窗口左上角移动到计算出的位置（确保窗口在屏幕可视区域内）
        self.move(window_geometry.topLeft())
    
    def init_ui(self):
        """
        初始化用户界面
        """
        # 设置窗口标题和大小
        self.setWindowTitle("DICOM到NIfTI转换器")
        self.setGeometry(100, 100, 500, 350)
        # 在显示屏中央显示
        self.center_on_screen()
        
        # 创建中心部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QVBoxLayout(central_widget)
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(20, 20, 20, 20)
        
        # 添加说明标签
        info_label = QLabel("本工具用于将DICOM文件转换为NIfTI格式。请选择转换类型并设置相应参数。")
        info_label.setWordWrap(True)
        main_layout.addWidget(info_label)
        
        # 创建转换类型选择区域
        conversion_type_group = QGroupBox("转换类型")
        conversion_type_layout = QVBoxLayout()  # 改为垂直布局以容纳三个选项
        
        self.type_group = QButtonGroup()
        self.dir_radio = QRadioButton("目录转换（多个系列到多个文件）")
        self.series_radio = QRadioButton("系列转换（单个系列到单个文件）")
        self.files_radio = QRadioButton("文件转换（多个文件到单个文件）")
        
        # 默认选择目录转换
        self.dir_radio.setChecked(True)
        self.dir_radio.toggled.connect(self.on_conversion_type_changed)
        self.series_radio.toggled.connect(self.on_conversion_type_changed)
        self.files_radio.toggled.connect(self.on_conversion_type_changed)
        
        self.type_group.addButton(self.dir_radio)
        self.type_group.addButton(self.series_radio)
        self.type_group.addButton(self.files_radio)
        
        conversion_type_layout.addWidget(self.dir_radio)
        conversion_type_layout.addWidget(self.series_radio)
        conversion_type_layout.addWidget(self.files_radio)
        conversion_type_group.setLayout(conversion_type_layout)
        main_layout.addWidget(conversion_type_group)
        
        # 创建输入路径选择区域
        input_layout = QHBoxLayout()
        self.input_label = QLabel("DICOM目录:")  # 改为实例变量以便修改
        self.input_path = QLineEdit()
        self.input_path.setReadOnly(True)
        self.browse_input_btn = QPushButton("浏览...")
        self.browse_input_btn.clicked.connect(self.browse_input)
        
        input_layout.addWidget(self.input_label)
        input_layout.addWidget(self.input_path, 1)
        input_layout.addWidget(self.browse_input_btn)
        main_layout.addLayout(input_layout)
        
        # 存储选中的DICOM文件列表
        self.selected_files = []
        
        # 创建输出路径选择区域
        output_layout = QHBoxLayout()
        self.output_label = QLabel("输出目录:")
        self.output_path = QLineEdit()
        self.output_path.setReadOnly(True)
        browse_output_btn = QPushButton("浏览...")
        browse_output_btn.clicked.connect(self.browse_output)
        
        output_layout.addWidget(self.output_label)
        output_layout.addWidget(self.output_path, 1)
        output_layout.addWidget(browse_output_btn)
        main_layout.addLayout(output_layout)
        
        # 创建选项设置区域
        options_group = QGroupBox("选项")
        options_layout = QHBoxLayout()
        
        self.reorient_check = QRadioButton("重定向NIfTI文件")
        self.reorient_check.setChecked(True)
        
        self.compression_check = QRadioButton("压缩NIfTI文件")
        self.compression_check.setChecked(True)
        
        options_layout.addWidget(self.reorient_check)
        options_layout.addWidget(self.compression_check)
        options_group.setLayout(options_layout)
        main_layout.addWidget(options_group)
        
        # 创建日志显示区域
        log_group = QGroupBox("日志输出")
        log_layout = QVBoxLayout()
        
        self.log_text_edit = QTextEdit()
        self.log_text_edit.setReadOnly(True)
        self.log_text_edit.setLineWrapMode(QTextEdit.WidgetWidth)
        
        # 创建滚动区域并添加文本编辑器
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setWidget(self.log_text_edit)
        
        log_layout.addWidget(scroll_area)
        log_group.setLayout(log_layout)
        log_group.setMinimumHeight(150)
        
        main_layout.addWidget(log_group)
        
        # 添加垂直拉伸
        main_layout.addStretch()
        
        # 创建转换按钮
        button_layout = QHBoxLayout()
        button_layout.addStretch()
        
        self.convert_btn = QPushButton("开始转换")
        self.convert_btn.setMinimumWidth(120)
        self.convert_btn.clicked.connect(self.start_conversion)
        
        self.cancel_btn = QPushButton("取消")
        self.cancel_btn.setMinimumWidth(120)
        self.cancel_btn.clicked.connect(self.cancel_conversion)
        
        button_layout.addWidget(self.convert_btn)
        button_layout.addWidget(self.cancel_btn)
        main_layout.addLayout(button_layout)
    
    def on_conversion_type_changed(self):
        """
        当转换类型改变时更新界面
        """
        if self.dir_radio.isChecked():
            self.input_label.setText("DICOM目录:")
            self.output_label.setText("输出目录:")
            self.compression_check.setVisible(True)
        elif self.series_radio.isChecked():
            self.input_label.setText("DICOM目录:")
            self.output_label.setText("输出文件:")
            self.compression_check.setVisible(False)
        else:  # 文件转换
            self.input_label.setText("DICOM文件:")
            self.output_label.setText("输出文件:")
            self.compression_check.setVisible(False)
    
    def browse_input(self):
        """
        根据转换类型浏览并选择输入（目录或文件）
        """
        if self.files_radio.isChecked():
            # 文件转换：选择多个DICOM文件
            files, _ = QFileDialog.getOpenFileNames(
                self, 
                "选择DICOM文件", 
                "", 
                "DICOM文件 (*.dcm);;所有文件 (*.*)"
            )
            if files:
                self.selected_files = files
                # 在输入框中显示文件数量和第一个文件路径
                if len(files) == 1:
                    self.input_path.setText(files[0])
                else:
                    # 显示第一个文件路径和总文件数
                    self.input_path.setText(f"{files[0]} 等共 {len(files)} 个文件")
        else:
            # 目录或系列转换：选择目录
            directory = QFileDialog.getExistingDirectory(
                self, 
                "选择DICOM目录",
                "",
                QFileDialog.ShowDirsOnly
            )
            if directory:
                self.selected_files = []  # 清空文件列表
                self.input_path.setText(directory)
    
    def browse_output(self):
        """
        浏览并选择输出路径（目录或文件）
        """
        if self.dir_radio.isChecked():
            # 目录转换：选择输出目录
            directory = QFileDialog.getExistingDirectory(self, "选择输出目录")
            if directory:
                self.output_path.setText(directory)
        else:
            # 系列转换：选择输出文件
            file_path, _ = QFileDialog.getSaveFileName(
                self, "保存NIfTI文件", "", "NIfTI文件 (*.nii.gz *.nii)"
            )
            if file_path:
                # 确保文件扩展名正确
                if not (file_path.endswith('.nii.gz') or file_path.endswith('.nii')):
                    file_path += '.nii.gz'
                self.output_path.setText(file_path)
    
    def start_conversion(self):
        """
        开始DICOM到NIfTI的转换过程
        """
        # 检查dicom2nifti库是否可用
        if not DICOM2NIFTI_AVAILABLE:
            QMessageBox.critical(self, "错误", "未找到dicom2nifti库，请使用pip install dicom2nifti进行安装")
            return
        
        # 初始化日志系统（仅在第一次转换时执行）
        if not self.logging_initialized:
            self.setup_logging()
            # 清空日志显示区域
            self.log_text_edit.clear()
            # 启动日志定时器
            self.log_timer.start(500)  # 每500毫秒刷新一次
            # 记录开始转换的日志
            logging.info("开始转换任务...")
        
        # 根据转换类型检查输入
        if self.files_radio.isChecked():
            # 文件转换：检查是否选择了文件
            if not self.selected_files:
                QMessageBox.warning(self, "警告", "请选择至少一个DICOM文件")
                return
            # 验证所有文件是否存在
            for file_path in self.selected_files:
                if not os.path.isfile(file_path):
                    QMessageBox.warning(self, "警告", f"文件不存在: {file_path}")
                    return
        else:
            # 目录或系列转换：检查目录
            dicom_path = self.input_path.text().strip()
            if not dicom_path or not os.path.isdir(dicom_path):
                QMessageBox.warning(self, "警告", "请选择有效的DICOM目录")
                return
        
        # 检查输出路径
        output_path = self.output_path.text().strip()
        if not output_path:
            QMessageBox.warning(self, "警告", "请选择输出路径")
            return
        
        # 根据转换类型处理输出路径
        if self.dir_radio.isChecked():
            # 目录转换：确保输出目录存在
            if not os.path.exists(output_path):
                try:
                    os.makedirs(output_path)
                except Exception as e:
                    QMessageBox.critical(self, "错误", f"无法创建输出目录: {str(e)}")
                    return
        else:
            # 系列或文件转换：确保输出文件的父目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                try:
                    os.makedirs(output_dir)
                except Exception as e:
                    QMessageBox.critical(self, "错误", f"无法创建输出目录: {str(e)}")
                    return
            # 确保文件扩展名正确
            if not (output_path.endswith('.nii.gz') or output_path.endswith('.nii')):
                output_path += '.nii.gz'
        
        # 获取选项设置
        reorient = self.reorient_check.isChecked()
        compression = self.compression_check.isChecked()
        
        # 禁用按钮
        self.convert_btn.setEnabled(False)
        
        # 移除进度条显示
        
        # 创建并启动转换工作线程
        if self.dir_radio.isChecked():
            conversion_type = 'directory'
            dicom_path = self.input_path.text().strip()
            self.worker = ConversionWorker(
                conversion_type, 
                dicom_path, 
                output_path, 
                reorient, 
                compression
            )
        elif self.series_radio.isChecked():
            conversion_type = 'series'
            dicom_path = self.input_path.text().strip()
            self.worker = ConversionWorker(
                conversion_type, 
                dicom_path, 
                output_path, 
                reorient, 
                compression
            )
        else:  # 文件转换
            conversion_type = 'files'
            # 对于文件转换，我们需要传递文件列表
            # 使用第一个文件的目录作为dicom_path（虽然实际上不会用到）
            dicom_path = os.path.dirname(self.selected_files[0]) if self.selected_files else ''
            self.worker = ConversionWorker(
                conversion_type, 
                dicom_path, 
                output_path, 
                reorient, 
                compression,
                self.selected_files
            )
        
        self.worker.finished.connect(self.on_conversion_finished)
        self.worker.start()
    
    # 移除进度条更新方法
    
    def cancel_conversion(self):
        """
        取消正在进行的转换操作，直接执行取消而不显示确认弹窗
        没有转换任务时不执行任何操作，不再关闭应用
        """
        # 检查是否有正在运行的转换
        if hasattr(self, 'worker') and self.worker and self.worker.isRunning():
            # 直接终止工作线程，无需确认
            self.worker.terminate()
            # 等待线程结束
            self.worker.wait()
            
            # 重置UI状态
            self.convert_btn.setEnabled(True)
            
            # 只记录日志，不显示弹窗
            logging.info("用户取消了转换操作")
        # 移除else分支，不再关闭应用
            
    def on_conversion_finished(self, success, message):
        """
        当转换完成时处理结果
        
        Args:
            success: 转换是否成功
            message: 结果消息
        """
        # 启用按钮
        self.convert_btn.setEnabled(True)
        
        # 记录并显示结果消息
        if success:
            logging.info(message)
        else:
            logging.error(message)
        
        # 记录转换完成日志
        logging.info("转换任务完成")
    
    def closeEvent(self, event):
        """
        窗口关闭时停止日志定时器
        """
        if hasattr(self, 'log_timer'):
            self.log_timer.stop()
        event.accept()


def main():
    """
    应用程序入口函数
    """
    # 创建应用程序实例
    app = QApplication(sys.argv)
    
    # 设置应用程序样式
    app.setStyle("Fusion")
    
    # 创建并显示主窗口
    window = Dicom2NiftiConverter()
    window.show()
    
    # 运行应用程序
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()